{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tHktDmcxoQBC"
   },
   "outputs": [],
   "source": [
    "!pip install transformers datasets tokenizers\n",
    "!wget http://www.cs.cornell.edu/~cristian/data/cornell_movie_dialogs_corpus.zip\n",
    "!unzip -qq cornell_movie_dialogs_corpus.zip\n",
    "!rm cornell_movie_dialogs_corpus.zip\n",
    "!mkdir datasets\n",
    "!mv cornell\\ movie-dialogs\\ corpus/movie_conversations.txt ./datasets\n",
    "!mv cornell\\ movie-dialogs\\ corpus/movie_lines.txt ./datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "movie_conversations.txt\n",
    "\n",
    "character1ID +++$+++ character2ID +++$+++ movieID +++$+++ conversationIDList\n",
    "\n",
    "u0 +++$+++ u2 +++$+++ m0 +++$+++ ['L194', 'L195', 'L196', 'L197']\n",
    "his example indicates that the conversation involves two characters (u0 and u2) in a movie (m0), and the conversation consists of the lines with IDs L194, L195, L196, and L197.\n",
    "\n",
    "\n",
    "movie_lines.txt\n",
    "\n",
    "lineID +++$+++ characterID +++$+++ movieID +++$+++ characterName +++$+++ text\n",
    "\n",
    "L194 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ Can we make this quick? Roxanne Korrine and Andrew Barrett are having an incredibly horrendous public break-up on the quad. Again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "思路：\n",
    "1、解析数据，形成pairs，pairs[0]为前一个句子，pairs[1]为后一个句子\n",
    "2、数据集加载和生成训练数据，根据规则，0.5为上下句，0.5不是，0.15的为masked（其中80%为[mask],10%为其他的随机的token\n",
    "3、"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zsvxdkMmPMsg"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import torch\n",
    "import re\n",
    "import random\n",
    "import transformers, datasets\n",
    "from tokenizers import BertWordPieceTokenizer\n",
    "from transformers import BertTokenizer\n",
    "import tqdm\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import itertools\n",
    "import math\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from torch.optim import Adam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DyRDEfhtoDRk"
   },
   "source": [
    "# 1 ) Tokenization (Word Piece Tokenizer)\n",
    "\n",
    "[Huggingface WordPieceTokenizer](https://huggingface.co/learn/nlp-course/chapter6/6?fw=pt)\n",
    "\n",
    "The tokenizer's primary job is to split the input text into smaller tokens. These tokens are usually words, subwords (WordPiece tokens), or characters, depending on the specific tokenizer and its configuration.\n",
    "\n",
    "Subword Tokenization (WordPiece): BERT often uses subword tokenization, where words are further divided into smaller units called subword tokens. For instance, \"unhappiness\" might be broken down into [\"un\", \"##hap\", \"##piness\"]\n",
    "\n",
    "\n",
    "By dividing the frequency of the pair by the product of the frequencies of each of its parts, the algorithm prioritizes the merging of pairs where the individual parts are less frequent in the vocabulary.\n",
    "\n",
    "**score=(freq_of_pair)/(freq_of_first_element×freq_of_second_element)**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kKAlXUvhxC5Y"
   },
   "source": [
    "## 1.1 Tokenizer from Scratch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 615,
     "status": "ok",
     "timestamp": 1702700487200,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "lvotJZcyxMai",
    "outputId": "0ac5ef7c-98d0-41cb-9af6-16ae03ab44e9"
   },
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n",
    "\n",
    "corpus = [\n",
    "    \"This is the Hugging Face Course.\",\n",
    "    \"This chapter is about tokenization.\",\n",
    "    \"This section shows several tokenizer algorithms.\",\n",
    "    \"Hopefully, you will be able to understand how they are trained and generate tokens.\",\n",
    "]\n",
    "\n",
    "### get the frequency of each word ###\n",
    "word_freqs = defaultdict(int)\n",
    "for text in corpus:\n",
    "    words_with_offsets = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(text)\n",
    "    new_words = [word for word, offset in words_with_offsets]\n",
    "    print(new_words)\n",
    "    for word in new_words:\n",
    "        word_freqs[word] += 1\n",
    "\n",
    "print(f\"\\nFinal Word Frequency: {word_freqs}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 522,
     "status": "ok",
     "timestamp": 1702700504613,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "FIHFIJwpxcbD",
    "outputId": "d0f83212-f6b0-4405-dc3f-c4b4cc64f039"
   },
   "outputs": [],
   "source": [
    "### split all word into alphabet ###\n",
    "alphabet = []\n",
    "for word in word_freqs.keys():\n",
    "    if word[0] not in alphabet:\n",
    "        alphabet.append(word[0])\n",
    "    for letter in word[1:]:\n",
    "        if f\"##{letter}\" not in alphabet:\n",
    "            alphabet.append(f\"##{letter}\")\n",
    "\n",
    "alphabet.sort()\n",
    "print(f'All alphabets: {alphabet}')\n",
    "\n",
    "### insert special token and subword ###\n",
    "vocab = [\"[PAD]\", \"[UNK]\", \"[CLS]\", \"[SEP]\", \"[MASK]\"] + alphabet.copy()\n",
    "splits = {word: [c if i == 0 else f\"##{c}\" for i, c in enumerate(word)] for word in word_freqs.keys()}\n",
    "print(f'\\nSplitted Words: {splits}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 346,
     "status": "ok",
     "timestamp": 1702700508853,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "M5svp-f028FO",
    "outputId": "9af94a71-8447-411e-8449-675050eb69ac"
   },
   "outputs": [],
   "source": [
    " ### compute score for merging ###\n",
    "def compute_pair_scores(splits):\n",
    "    letter_freqs = defaultdict(int)\n",
    "    pair_freqs = defaultdict(int)\n",
    "\n",
    "    for word, freq in word_freqs.items():\n",
    "        split = splits[word]\n",
    "        if len(split) == 1:\n",
    "            letter_freqs[split[0]] += freq\n",
    "            continue\n",
    "        for i in range(len(split) - 1):\n",
    "            pair = (split[i], split[i + 1])\n",
    "            letter_freqs[split[i]] += freq\n",
    "            pair_freqs[pair] += freq\n",
    "        letter_freqs[split[-1]] += freq\n",
    "\n",
    "    scores = {\n",
    "        pair: freq / (letter_freqs[pair[0]] * letter_freqs[pair[1]])\n",
    "        for pair, freq in pair_freqs.items()\n",
    "    }\n",
    "    return scores\n",
    "\n",
    "pair_scores = compute_pair_scores(splits)\n",
    "print(f'Scores for each Pair: {pair_scores}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 327,
     "status": "ok",
     "timestamp": 1702700639209,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "dkcQKo1r2-c5",
    "outputId": "40267815-b5c2-4148-cdbc-8a8cb211cf65"
   },
   "outputs": [],
   "source": [
    "### finding pair with best score ###\n",
    "best_pair = \"\"\n",
    "max_score = None\n",
    "for pair, score in pair_scores.items():\n",
    "    if max_score is None or max_score < score:\n",
    "        best_pair = pair\n",
    "        max_score = score\n",
    "\n",
    "print(best_pair, max_score)\n",
    "vocab.append(\"ab\")\n",
    "\n",
    "### merge pair ###\n",
    "def merge_pair(a, b, splits):\n",
    "    for word in word_freqs:\n",
    "        split = splits[word]\n",
    "        if len(split) == 1:\n",
    "            continue\n",
    "        i = 0\n",
    "        while i < len(split) - 1:\n",
    "            if split[i] == a and split[i + 1] == b:\n",
    "                merge = a + b[2:] if b.startswith(\"##\") else a + b\n",
    "                split = split[:i] + [merge] + split[i + 2 :]\n",
    "            else:\n",
    "                i += 1\n",
    "        splits[word] = split\n",
    "    return splits\n",
    "\n",
    "splits = merge_pair(\"a\", \"##b\", splits)\n",
    "print(splits[\"about\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 350,
     "status": "ok",
     "timestamp": 1702700824732,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "bc6S6Tj93sh8",
    "outputId": "1afabfb5-50ce-4752-f07f-2776d62a5c4e"
   },
   "outputs": [],
   "source": [
    "### keep looping to merge more pair\n",
    "vocab_size = 70\n",
    "while len(vocab) < vocab_size:\n",
    "    scores = compute_pair_scores(splits)\n",
    "    best_pair, max_score = \"\", None\n",
    "    for pair, score in scores.items():\n",
    "        if max_score is None or max_score < score:\n",
    "            best_pair = pair\n",
    "            max_score = score\n",
    "    splits = merge_pair(*best_pair, splits)\n",
    "    new_token = (\n",
    "        best_pair[0] + best_pair[1][2:]\n",
    "        if best_pair[1].startswith(\"##\")\n",
    "        else best_pair[0] + best_pair[1]\n",
    "    )\n",
    "    vocab.append(new_token)\n",
    "\n",
    "print(f'Final Vocab: {vocab}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 357,
     "status": "ok",
     "timestamp": 1702700889155,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "vFVyHHe84Vbi",
    "outputId": "afb79be0-eda1-40a3-8a77-dc2f2078019b"
   },
   "outputs": [],
   "source": [
    "### ro encode a word ###\n",
    "def encode_word(word):\n",
    "    tokens = []\n",
    "    while len(word) > 0:\n",
    "        i = len(word)\n",
    "        while i > 0 and word[:i] not in vocab:\n",
    "            i -= 1\n",
    "        if i == 0:\n",
    "            return [\"[UNK]\"]\n",
    "        tokens.append(word[:i])\n",
    "        word = word[i:]\n",
    "        if len(word) > 0:\n",
    "            word = f\"##{word}\"\n",
    "    return tokens\n",
    "\n",
    "print(encode_word(\"Hugging\"))\n",
    "print(encode_word(\"HOgging\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5jaQsRw4xC-x"
   },
   "source": [
    "## 1.2 Tokenizer Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "movie_conversations.txt\n",
    "\n",
    "character1ID +++$+++ character2ID +++$+++ movieID +++$+++ conversationIDList\n",
    "\n",
    "u0 +++$+++ u2 +++$+++ m0 +++$+++ ['L194', 'L195', 'L196', 'L197']\n",
    "his example indicates that the conversation involves two characters (u0 and u2) in a movie (m0), and the conversation consists of the lines with IDs L194, L195, L196, and L197.\n",
    "\n",
    "\n",
    "movie_lines.txt\n",
    "\n",
    "lineID +++$+++ characterID +++$+++ movieID +++$+++ characterName +++$+++ text\n",
    "\n",
    "L194 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ Can we make this quick? Roxanne Korrine and Andrew Barrett are having an incredibly horrendous public break-up on the quad. Again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "关于电影对话语料库的内容，可以从对话中找出上下相邻的句子，这里取对话的前64个token（word）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1961,
     "status": "ok",
     "timestamp": 1665133698134,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "ER4dwYJDoZFU",
    "outputId": "c0ca8f0c-4610-4dc2-f386-09be4d2f139f"
   },
   "outputs": [],
   "source": [
    "### data processing\n",
    "MAX_LEN = 64\n",
    "\n",
    "### loading all data into memory\n",
    "corpus_movie_conv = './datasets/movie_conversations.txt'\n",
    "corpus_movie_lines = './datasets/movie_lines.txt'\n",
    "with open(corpus_movie_conv, 'r', encoding='iso-8859-1') as c:\n",
    "    conv = c.readlines()\n",
    "with open(corpus_movie_lines, 'r', encoding='iso-8859-1') as l:\n",
    "    lines = l.readlines()\n",
    "\n",
    "### splitting text using special lines\n",
    "lines_dic = {}\n",
    "for line in lines:\n",
    "    objects = line.split(\" +++$+++ \")\n",
    "    lines_dic[objects[0]] = objects[-1]\n",
    "\n",
    "### generate question answer pairs\n",
    "pairs = []\n",
    "for con in conv:\n",
    "    ids = eval(con.split(\" +++$+++ \")[-1])  #获取一个对话的行数列表，列表中相邻的行即为上下临近的句子，这里取对话的前64个token（word）\n",
    "    for i in range(len(ids)):\n",
    "        qa_pairs = []\n",
    "\n",
    "        if i == len(ids) - 1:\n",
    "            break\n",
    "\n",
    "        first = lines_dic[ids[i]].strip()\n",
    "        second = lines_dic[ids[i+1]].strip()\n",
    "\n",
    "        qa_pairs.append(' '.join(first.split()[:MAX_LEN])) #split对字符串切分成list，如果split没有给定参数，就用默认参数（空格）\n",
    "        qa_pairs.append(' '.join(second.split()[:MAX_LEN]))\n",
    "        pairs.append(qa_pairs)\n",
    "\n",
    "# sample\n",
    "print(pairs[20])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 868,
     "status": "ok",
     "timestamp": 1665133759805,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "c4K2nIFalnOz",
    "outputId": "9cda79b0-7da1-467a-ba83-09f1f0f9095d"
   },
   "outputs": [],
   "source": [
    "# WordPiece tokenizer\n",
    "### save data as txt file\n",
    "os.mkdir('./data')\n",
    "text_data = []\n",
    "file_count = 0\n",
    "\n",
    "for sample in tqdm.tqdm([x[0] for x in pairs]):\n",
    "    text_data.append(sample)\n",
    "\n",
    "    # once we hit the 10K mark, save to file\n",
    "    if len(text_data) == 10000:\n",
    "        with open(f'./data/text_{file_count}.txt', 'w', encoding='utf-8') as fp:\n",
    "            fp.write('\\n'.join(text_data))\n",
    "        text_data = []\n",
    "        file_count += 1\n",
    "\n",
    "paths = [str(x) for x in Path('./data').glob('**/*.txt')]\n",
    "print(len(paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 3975,
     "status": "ok",
     "timestamp": 1665133790644,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "TzMa6PMWrJZj",
    "outputId": "9ba1f2ff-c7c5-42a3-a4d3-5de0ea35b338"
   },
   "outputs": [],
   "source": [
    "### training own tokenizer\n",
    "tokenizer = BertWordPieceTokenizer(\n",
    "    clean_text=True,\n",
    "    handle_chinese_chars=False,\n",
    "    strip_accents=False,\n",
    "    lowercase=True\n",
    ")\n",
    "\n",
    "tokenizer.train(\n",
    "    files=paths,\n",
    "    vocab_size=30_000,\n",
    "    min_frequency=5,\n",
    "    limit_alphabet=1000,\n",
    "    wordpieces_prefix='##',\n",
    "    special_tokens=['[PAD]', '[CLS]', '[SEP]', '[MASK]', '[UNK]']\n",
    "    )\n",
    "\n",
    "os.mkdir('./bert-it-1')\n",
    "tokenizer.save_model('./bert-it-1', 'bert-it')\n",
    "tokenizer = BertTokenizer.from_pretrained('./bert-it-1/bert-it-vocab.txt', local_files_only=True)\n",
    "token_ids = tokenizer('I like surfboarding!')['input_ids']\n",
    "print(token_ids)\n",
    "print(tokenizer.convert_ids_to_tokens(token_ids))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sbBU9BTdvlQ0"
   },
   "source": [
    "# 2) Pre-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = [1,2,3]\n",
    "b = a[:5]\n",
    "\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bXCsfR3tmajw"
   },
   "outputs": [],
   "source": [
    "class BERTDataset(Dataset):\n",
    "    def __init__(self, data_pair, tokenizer, seq_len=64):\n",
    "\n",
    "        self.tokenizer = tokenizer\n",
    "        self.seq_len = seq_len\n",
    "        self.corpus_lines = len(data_pair)\n",
    "        self.lines = data_pair\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.corpus_lines\n",
    "\n",
    "    def __getitem__(self, item):\n",
    "\n",
    "        # Step 1: get random sentence pair, either negative or positive (saved as is_next_label)\n",
    "        t1, t2, is_next_label = self.get_sent(item)\n",
    "\n",
    "        # Step 2: replace random words in sentence with mask / random words\n",
    "        t1_random, t1_label = self.random_word(t1)\n",
    "        t2_random, t2_label = self.random_word(t2)\n",
    "\n",
    "        # Step 3: Adding CLS and SEP tokens to the start and end of sentences\n",
    "        # Adding PAD token for labels \n",
    "        # 输出的label将[CLS]替换为了PAD，在MLM中将不参与loss的计算\n",
    "        t1 = [self.tokenizer.vocab['[CLS]']] + t1_random + [self.tokenizer.vocab['[SEP]']]\n",
    "        t2 = t2_random + [self.tokenizer.vocab['[SEP]']]\n",
    "        t1_label = [self.tokenizer.vocab['[PAD]']] + t1_label + [self.tokenizer.vocab['[PAD]']]\n",
    "        t2_label = t2_label + [self.tokenizer.vocab['[PAD]']]\n",
    "\n",
    "        # Step 4: combine sentence 1 and 2 as one input\n",
    "        # adding PAD tokens to make the sentence same length as seq_len\n",
    "        segment_label = ([1 for _ in range(len(t1))] + [2 for _ in range(len(t2))])[:self.seq_len]\n",
    "        bert_input = (t1 + t2)[:self.seq_len]\n",
    "        bert_label = (t1_label + t2_label)[:self.seq_len]\n",
    "        padding = [self.tokenizer.vocab['[PAD]'] for _ in range(self.seq_len - len(bert_input))]\n",
    "        bert_input.extend(padding), bert_label.extend(padding), segment_label.extend(padding)\n",
    "\n",
    "        output = {\"bert_input\": bert_input,\n",
    "                  \"bert_label\": bert_label,\n",
    "                  \"segment_label\": segment_label,\n",
    "                  \"is_next\": is_next_label}\n",
    "\n",
    "        return {key: torch.tensor(value) for key, value in output.items()}\n",
    "\n",
    "    def random_word(self, sentence):\n",
    "        tokens = sentence.split()\n",
    "        output_label = []\n",
    "        output = []\n",
    "\n",
    "        # 15% of the tokens would be replaced\n",
    "        for i, token in enumerate(tokens):\n",
    "            prob = random.random()\n",
    "\n",
    "            # remove cls and sep token\n",
    "            token_id = self.tokenizer(token)['input_ids'][1:-1] #有可能生成多个subword\n",
    "\n",
    "            # 15% chance of altering token\n",
    "            if prob < 0.15:\n",
    "                prob /= 0.15\n",
    "\n",
    "                # 80% chance change token to mask token\n",
    "                if prob < 0.8:\n",
    "                    for i in range(len(token_id)):\n",
    "                        output.append(self.tokenizer.vocab['[MASK]'])\n",
    "\n",
    "                # 10% chance change token to random token\n",
    "                elif prob < 0.9:\n",
    "                    for i in range(len(token_id)):\n",
    "                        output.append(random.randrange(len(self.tokenizer.vocab)))\n",
    "\n",
    "                # 10% chance change token to current token\n",
    "                else:\n",
    "                    output.append(token_id)\n",
    "\n",
    "                output_label.append(token_id)#真值保持不变\n",
    "\n",
    "            else:\n",
    "                output.append(token_id)\n",
    "                for i in range(len(token_id)):\n",
    "                    output_label.append(0)#真值为0，不参与loss的计算\n",
    "\n",
    "        # flattening\n",
    "        output = list(itertools.chain(*[[x] if not isinstance(x, list) else x for x in output]))\n",
    "        output_label = list(itertools.chain(*[[x] if not isinstance(x, list) else x for x in output_label]))\n",
    "        assert len(output) == len(output_label)\n",
    "        return output, output_label\n",
    "\n",
    "    def get_sent(self, index):\n",
    "        '''return random sentence pair'''\n",
    "        t1, t2 = self.get_corpus_line(index)\n",
    "\n",
    "        # negative or positive pair, for next sentence prediction\n",
    "        if random.random() > 0.5:\n",
    "            return t1, t2, 1\n",
    "        else:\n",
    "            return t1, self.get_random_line(), 0\n",
    "\n",
    "    def get_corpus_line(self, item):\n",
    "        '''return sentence pair'''\n",
    "        return self.lines[item][0], self.lines[item][1]\n",
    "\n",
    "    def get_random_line(self):\n",
    "        '''return random single sentence'''\n",
    "        return self.lines[random.randrange(len(self.lines))][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4637,
     "status": "ok",
     "timestamp": 1665133798321,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "5RD5ta6KA_wK",
    "outputId": "3b408feb-a0ff-4d66-89a7-6e7e32a57ee3"
   },
   "outputs": [],
   "source": [
    "# test\n",
    "print(\"\\n\")\n",
    "train_data = BERTDataset(pairs, seq_len=MAX_LEN, tokenizer=tokenizer)\n",
    "train_loader = DataLoader(train_data, batch_size=32, shuffle=True, pin_memory=True)\n",
    "sample_data = next(iter(train_loader))\n",
    "print('Batch Size', sample_data['bert_input'].size())\n",
    "\n",
    "# 3 is MASK\n",
    "result = train_data[random.randrange(len(train_data))]\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AT-uR5BXChcM"
   },
   "source": [
    "# 3) Modeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 804,
     "status": "ok",
     "timestamp": 1665133801248,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "S6Nr8gMt49dF",
    "outputId": "e834d002-629f-4d93-a5dd-4e56bf07ec1f"
   },
   "outputs": [],
   "source": [
    "### embedding\n",
    "class PositionalEmbedding(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, max_len=128):\n",
    "        super().__init__()\n",
    "\n",
    "        # Compute the positional encodings once in log space.\n",
    "        pe = torch.zeros(max_len, d_model).float()\n",
    "        pe.require_grad = False\n",
    "\n",
    "        for pos in range(max_len):\n",
    "            # for each dimension of the each position\n",
    "            for i in range(0, d_model, 2):\n",
    "                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i)/d_model)))\n",
    "                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1))/d_model)))\n",
    "\n",
    "        # include the batch size\n",
    "        self.pe = pe.unsqueeze(0)\n",
    "        # self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.pe\n",
    "\n",
    "class BERTEmbedding(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    BERT Embedding which is consisted with under features\n",
    "        1. TokenEmbedding : normal embedding matrix\n",
    "        2. PositionalEmbedding : adding positional information using sin, cos\n",
    "        2. SegmentEmbedding : adding sentence segment info, (sent_A:1, sent_B:2)\n",
    "        sum of all these features are output of BERTEmbedding\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, embed_size, seq_len=64, dropout=0.1):\n",
    "        \"\"\"\n",
    "        :param vocab_size: total vocab size\n",
    "        :param embed_size: embedding size of token embedding\n",
    "        :param dropout: dropout rate\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.embed_size = embed_size\n",
    "        # (m, seq_len) --> (m, seq_len, embed_size)\n",
    "        # padding_idx is not updated during training, remains as fixed pad (0)\n",
    "        self.token = torch.nn.Embedding(vocab_size, embed_size, padding_idx=0)\n",
    "        self.segment = torch.nn.Embedding(3, embed_size, padding_idx=0)\n",
    "        self.position = PositionalEmbedding(d_model=embed_size, max_len=seq_len)\n",
    "        self.dropout = torch.nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, sequence, segment_label):\n",
    "        x = self.token(sequence) + self.position(sequence) + self.segment(segment_label)\n",
    "        return self.dropout(x)\n",
    "\n",
    "### testing\n",
    "embed_layer = BERTEmbedding(vocab_size=len(tokenizer.vocab), embed_size=768, seq_len=MAX_LEN)\n",
    "embed_result = embed_layer(sample_data['bert_input'], sample_data['segment_label'])\n",
    "print(embed_result.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1036,
     "status": "ok",
     "timestamp": 1665133802282,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "g7n6FOkOBWK2",
    "outputId": "61857fc3-c499-4a4f-a97d-2d5af268e2b1"
   },
   "outputs": [],
   "source": [
    "### attention layers\n",
    "class MultiHeadedAttention(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, heads, d_model, dropout=0.1):\n",
    "        super(MultiHeadedAttention, self).__init__()\n",
    "\n",
    "        assert d_model % heads == 0\n",
    "        self.d_k = d_model // heads\n",
    "        self.heads = heads\n",
    "        self.dropout = torch.nn.Dropout(dropout)\n",
    "\n",
    "        self.query = torch.nn.Linear(d_model, d_model)\n",
    "        self.key = torch.nn.Linear(d_model, d_model)\n",
    "        self.value = torch.nn.Linear(d_model, d_model)\n",
    "        self.output_linear = torch.nn.Linear(d_model, d_model)\n",
    "\n",
    "    def forward(self, query, key, value, mask):\n",
    "        \"\"\"\n",
    "        query, key, value of shape: (batch_size, max_len, d_model)\n",
    "        mask of shape: (batch_size, 1, 1, max_words)\n",
    "        \"\"\"\n",
    "        # (batch_size, max_len, d_model)\n",
    "        query = self.query(query)\n",
    "        key = self.key(key)\n",
    "        value = self.value(value)\n",
    "\n",
    "        # (batch_size, max_len, d_model) --> (batch_size, max_len, h, d_k) --> (batch_size, h, max_len, d_k)\n",
    "        query = query.view(query.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)\n",
    "        key = key.view(key.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)\n",
    "        value = value.view(value.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)\n",
    "\n",
    "        # (batch_size, h, max_len, d_k) matmul (batch_size, h, d_k, max_len) --> (batch_size, h, max_len, max_len)\n",
    "        scores = torch.matmul(query, key.permute(0, 1, 3, 2)) / math.sqrt(query.size(-1))\n",
    "\n",
    "        # fill 0 mask with super small number so it wont affect the softmax weight\n",
    "        # (batch_size, h, max_len, max_len)\n",
    "        scores = scores.masked_fill(mask == 0, -1e9)\n",
    "\n",
    "        # (batch_size, h, max_len, max_len)\n",
    "        # softmax to put attention weight for all non-pad tokens\n",
    "        # max_len X max_len matrix of attention\n",
    "        weights = F.softmax(scores, dim=-1)\n",
    "        weights = self.dropout(weights)\n",
    "\n",
    "        # (batch_size, h, max_len, max_len) matmul (batch_size, h, max_len, d_k) --> (batch_size, h, max_len, d_k)\n",
    "        context = torch.matmul(weights, value)\n",
    "\n",
    "        # (batch_size, h, max_len, d_k) --> (batch_size, max_len, h, d_k) --> (batch_size, max_len, d_model)\n",
    "        context = context.permute(0, 2, 1, 3).contiguous().view(context.shape[0], -1, self.heads * self.d_k)\n",
    "\n",
    "        # (batch_size, max_len, d_model)\n",
    "        return self.output_linear(context)\n",
    "\n",
    "class FeedForward(torch.nn.Module):\n",
    "    \"Implements FFN equation\"\n",
    "\n",
    "    def __init__(self, d_model, middle_dim=2048, dropout=0.1):\n",
    "        super(FeedForward, self).__init__()\n",
    "\n",
    "        self.fc1 = torch.nn.Linear(d_model, middle_dim)\n",
    "        self.fc2 = torch.nn.Linear(middle_dim, d_model)\n",
    "        self.dropout = torch.nn.Dropout(dropout)\n",
    "        self.activation = torch.nn.GELU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.activation(self.fc1(x))\n",
    "        out = self.fc2(self.dropout(out))\n",
    "        return out\n",
    "\n",
    "class EncoderLayer(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        d_model=768,\n",
    "        heads=12,\n",
    "        feed_forward_hidden=768 * 4,\n",
    "        dropout=0.1\n",
    "        ):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.layernorm = torch.nn.LayerNorm(d_model)\n",
    "        self.self_multihead = MultiHeadedAttention(heads, d_model)\n",
    "        self.feed_forward = FeedForward(d_model, middle_dim=feed_forward_hidden)\n",
    "        self.dropout = torch.nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, embeddings, mask):\n",
    "        # embeddings: (batch_size, max_len, d_model)\n",
    "        # encoder mask: (batch_size, 1, 1, max_len)\n",
    "        # result: (batch_size, max_len, d_model)\n",
    "        interacted = self.dropout(self.self_multihead(embeddings, embeddings, embeddings, mask))\n",
    "        # residual layer\n",
    "        interacted = self.layernorm(interacted + embeddings)\n",
    "        # bottleneck\n",
    "        feed_forward_out = self.dropout(self.feed_forward(interacted))\n",
    "        encoded = self.layernorm(feed_forward_out + interacted)\n",
    "        return encoded\n",
    "\n",
    "### testing\n",
    "mask = (sample_data['bert_input'] > 0).unsqueeze(1).repeat(1, sample_data['bert_input'].size(1), 1).unsqueeze(1)\n",
    "transformer_block = EncoderLayer()\n",
    "transformer_result = transformer_block(embed_result, mask)\n",
    "transformer_result.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "mask = (sample_data['bert_input'] > 0).unsqueeze(1).repeat(1, sample_data['bert_input'].size(1), 1).unsqueeze(1)\n",
    "[batch_size, seq_len]-》[batch_size, 1, seq_len]-》[batch_size, seq_len, seq_len]-》[batch_size, 1, seq_len, seq_len]\n",
    "\n",
    "repeat函数的用法：the repeat function is used to repeat a tensor along specified dimensions. It effectively creates a new tensor by replicating the input tensor a specified number of times along each dimension.\n",
    "tensor.repeat(*sizes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 15539,
     "status": "ok",
     "timestamp": 1665133817819,
     "user": {
      "displayName": "Kean Chan",
      "userId": "05792587367281359063"
     },
     "user_tz": -480
    },
    "id": "T4am76N6Cimj",
    "outputId": "e0b849c5-50c2-4f15-b1f0-f08fbe95fc90"
   },
   "outputs": [],
   "source": [
    "class BERT(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    BERT model : Bidirectional Encoder Representations from Transformers.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, d_model=768, n_layers=12, heads=12, dropout=0.1):\n",
    "        \"\"\"\n",
    "        :param vocab_size: vocab_size of total words\n",
    "        :param hidden: BERT model hidden size\n",
    "        :param n_layers: numbers of Transformer blocks(layers)\n",
    "        :param attn_heads: number of attention heads\n",
    "        :param dropout: dropout rate\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.n_layers = n_layers\n",
    "        self.heads = heads\n",
    "\n",
    "        # paper noted they used 4*hidden_size for ff_network_hidden_size\n",
    "        self.feed_forward_hidden = d_model * 4\n",
    "\n",
    "        # embedding for BERT, sum of positional, segment, token embeddings\n",
    "        self.embedding = BERTEmbedding(vocab_size=vocab_size, embed_size=d_model)\n",
    "\n",
    "        # multi-layers transformer blocks, deep network\n",
    "        self.encoder_blocks = torch.nn.ModuleList(\n",
    "            [EncoderLayer(d_model, heads, d_model * 4, dropout) for _ in range(n_layers)])\n",
    "\n",
    "    def forward(self, x, segment_info):\n",
    "        # attention masking for padded token\n",
    "        # (batch_size, 1, seq_len, seq_len)\n",
    "        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)\n",
    "\n",
    "        # embedding the indexed sequence to sequence of vectors\n",
    "        x = self.embedding(x, segment_info)\n",
    "\n",
    "        # running over multiple transformer blocks\n",
    "        for encoder in self.encoder_blocks:\n",
    "            x = encoder.forward(x, mask)\n",
    "        return x\n",
    "\n",
    "class NextSentencePrediction(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    2-class classification model : is_next, is_not_next\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, hidden):\n",
    "        \"\"\"\n",
    "        :param hidden: BERT model output size\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.linear = torch.nn.Linear(hidden, 2)\n",
    "        self.softmax = torch.nn.LogSoftmax(dim=-1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # use only the first token which is the [CLS]\n",
    "        # [batch_size, seq_len, hidden_dim]\n",
    "        # x[:, 0] -> [batch_size, hidden_dim]\n",
    "        return self.softmax(self.linear(x[:, 0]))\n",
    "\n",
    "class MaskedLanguageModel(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    predicting origin token from masked input sequence\n",
    "    n-class classification problem, n-class = vocab_size\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, hidden, vocab_size):\n",
    "        \"\"\"\n",
    "        :param hidden: output size of BERT model\n",
    "        :param vocab_size: total vocab size\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.linear = torch.nn.Linear(hidden, vocab_size)\n",
    "        self.softmax = torch.nn.LogSoftmax(dim=-1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #[batch_size, seq_len, hidden_dim]\n",
    "        # 前面两个维度为批处理维度\n",
    "        return self.softmax(self.linear(x))\n",
    "\n",
    "class BERTLM(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    BERT Language Model\n",
    "    Next Sentence Prediction Model + Masked Language Model\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, bert: BERT, vocab_size):\n",
    "        \"\"\"\n",
    "        :param bert: BERT model which should be trained\n",
    "        :param vocab_size: total vocab size for masked_lm\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.bert = bert\n",
    "        self.next_sentence = NextSentencePrediction(self.bert.d_model)\n",
    "        self.mask_lm = MaskedLanguageModel(self.bert.d_model, vocab_size)\n",
    "\n",
    "    def forward(self, x, segment_label):\n",
    "        x = self.bert(x, segment_label)\n",
    "        return self.next_sentence(x), self.mask_lm(x)\n",
    "\n",
    "### test\n",
    "bert_model = BERT(len(tokenizer.vocab))\n",
    "bert_result = bert_model(sample_data['bert_input'], sample_data['segment_label'])\n",
    "print(bert_result.size())\n",
    "\n",
    "bert_lm = BERTLM(bert_model, len(tokenizer.vocab))\n",
    "final_result = bert_lm(sample_data['bert_input'], sample_data['segment_label'])\n",
    "print(final_result[0].size(), final_result[1].size())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nnp98JEZWwgN"
   },
   "source": [
    "# 4) Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0PPi4L1sCjBf"
   },
   "outputs": [],
   "source": [
    "### optimizer\n",
    "class ScheduledOptim():\n",
    "    '''A simple wrapper class for learning rate scheduling'''\n",
    "\n",
    "    def __init__(self, optimizer, d_model, n_warmup_steps):\n",
    "        self._optimizer = optimizer\n",
    "        self.n_warmup_steps = n_warmup_steps\n",
    "        self.n_current_steps = 0\n",
    "        self.init_lr = np.power(d_model, -0.5)\n",
    "\n",
    "    def step_and_update_lr(self):\n",
    "        \"Step with the inner optimizer\"\n",
    "        self._update_learning_rate()\n",
    "        self._optimizer.step()\n",
    "\n",
    "    def zero_grad(self):\n",
    "        \"Zero out the gradients by the inner optimizer\"\n",
    "        self._optimizer.zero_grad()\n",
    "\n",
    "    def _get_lr_scale(self):\n",
    "        return np.min([\n",
    "            np.power(self.n_current_steps, -0.5),\n",
    "            np.power(self.n_warmup_steps, -1.5) * self.n_current_steps])\n",
    "\n",
    "    def _update_learning_rate(self):\n",
    "        ''' Learning rate scheduling per step '''\n",
    "\n",
    "        self.n_current_steps += 1\n",
    "        lr = self.init_lr * self._get_lr_scale()\n",
    "\n",
    "        for param_group in self._optimizer.param_groups:\n",
    "            param_group['lr'] = lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IEggGpx2kUXx",
    "outputId": "578f0c3c-930b-41a8-8b3c-592d3888ecac"
   },
   "outputs": [],
   "source": [
    "### trainer\n",
    "class BERTTrainer:\n",
    "    def __init__(\n",
    "        self,\n",
    "        model,\n",
    "        train_dataloader,\n",
    "        test_dataloader=None,\n",
    "        lr= 1e-4,\n",
    "        weight_decay=0.01,\n",
    "        betas=(0.9, 0.999),\n",
    "        warmup_steps=10000,\n",
    "        log_freq=10,\n",
    "        device='cuda'\n",
    "        ):\n",
    "\n",
    "        self.device = device\n",
    "        self.model = model\n",
    "        self.train_data = train_dataloader\n",
    "        self.test_data = test_dataloader\n",
    "\n",
    "        # Setting the Adam optimizer with hyper-param\n",
    "        self.optim = Adam(self.model.parameters(), lr=lr, betas=betas, weight_decay=weight_decay)\n",
    "        self.optim_schedule = ScheduledOptim(\n",
    "            self.optim, self.model.bert.d_model, n_warmup_steps=warmup_steps\n",
    "            )\n",
    "\n",
    "        # Using Negative Log Likelihood Loss function for predicting the masked_token\n",
    "        self.criterion = torch.nn.NLLLoss(ignore_index=0)\n",
    "        self.log_freq = log_freq\n",
    "        print(\"Total Parameters:\", sum([p.nelement() for p in self.model.parameters()]))\n",
    "\n",
    "    def train(self, epoch):\n",
    "        self.iteration(epoch, self.train_data)\n",
    "\n",
    "    def test(self, epoch):\n",
    "        self.iteration(epoch, self.test_data, train=False)\n",
    "\n",
    "    def iteration(self, epoch, data_loader, train=True):\n",
    "\n",
    "        avg_loss = 0.0\n",
    "        total_correct = 0\n",
    "        total_element = 0\n",
    "\n",
    "        mode = \"train\" if train else \"test\"\n",
    "\n",
    "        # progress bar\n",
    "        data_iter = tqdm.tqdm(\n",
    "            enumerate(data_loader),\n",
    "            desc=\"EP_%s:%d\" % (mode, epoch),\n",
    "            total=len(data_loader),\n",
    "            bar_format=\"{l_bar}{r_bar}\"\n",
    "        )\n",
    "\n",
    "        for i, data in data_iter:\n",
    "\n",
    "            # 0. batch_data will be sent into the device(GPU or cpu)\n",
    "            data = {key: value.to(self.device) for key, value in data.items()}\n",
    "\n",
    "            # 1. forward the next_sentence_prediction and masked_lm model\n",
    "            # [batch_size,2], [batch_size, seq_len, vocab_size]\n",
    "            next_sent_output, mask_lm_output = self.model.forward(data[\"bert_input\"], data[\"segment_label\"])\n",
    "\n",
    "            # 2-1. NLL(negative log likelihood) loss of is_next classification result\n",
    "            next_loss = self.criterion(next_sent_output, data[\"is_next\"])\n",
    "\n",
    "            # 2-2. NLLLoss of predicting masked token word\n",
    "            # transpose to (m, vocab_size, seq_len) vs (m, seq_len)\n",
    "            # criterion(mask_lm_output.view(-1, mask_lm_output.size(-1)), data[\"bert_label\"].view(-1))\n",
    "            mask_loss = self.criterion(mask_lm_output.transpose(1, 2), data[\"bert_label\"])\n",
    "\n",
    "            # 2-3. Adding next_loss and mask_loss : 3.4 Pre-training Procedure\n",
    "            loss = next_loss + mask_loss\n",
    "\n",
    "            # 3. backward and optimization only in train\n",
    "            if train:\n",
    "                self.optim_schedule.zero_grad()\n",
    "                loss.backward()\n",
    "                self.optim_schedule.step_and_update_lr()\n",
    "\n",
    "            # next sentence prediction accuracy\n",
    "            correct = next_sent_output.argmax(dim=-1).eq(data[\"is_next\"]).sum().item()\n",
    "            avg_loss += loss.item()\n",
    "            total_correct += correct\n",
    "            total_element += data[\"is_next\"].nelement()\n",
    "\n",
    "            post_fix = {\n",
    "                \"epoch\": epoch,\n",
    "                \"iter\": i,\n",
    "                \"avg_loss\": avg_loss / (i + 1),\n",
    "                \"avg_acc\": total_correct / total_element * 100,\n",
    "                \"loss\": loss.item()\n",
    "            }\n",
    "\n",
    "            if i % self.log_freq == 0:\n",
    "                data_iter.write(str(post_fix))\n",
    "        print(\n",
    "            f\"EP{epoch}, {mode}: \\\n",
    "            avg_loss={avg_loss / len(data_iter)}, \\\n",
    "            total_acc={total_correct * 100.0 / total_element}\"\n",
    "        )\n",
    "\n",
    "### test\n",
    "train_data = BERTDataset(pairs, seq_len=MAX_LEN, tokenizer=tokenizer)\n",
    "train_loader = DataLoader(train_data, batch_size=32, shuffle=True, pin_memory=True)\n",
    "bert_model = BERT(len(tokenizer.vocab))\n",
    "bert_lm = BERTLM(bert_model, len(tokenizer.vocab))\n",
    "bert_trainer = BERTTrainer(bert_lm, train_loader, device='cpu')\n",
    "epochs = 2\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    bert_trainer.train(epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hvIpgHXg6HOK"
   },
   "source": [
    "# 5) Reference\n",
    "\n",
    "[BERT from Scratch](https://medium.com/data-and-beyond/complete-guide-to-building-bert-model-from-sratch-3e6562228891) | [BERT vs Roberta vs XLM](https://towardsdatascience.com/bert-roberta-distilbert-xlnet-which-one-to-use-3d5ab82ba5f8) | [StructBert vs Albert vs LongForm](https://towardsdatascience.com/advancing-over-bert-bigbird-convbert-dynabert-bca78a45629c) | [BART](https://medium.com/analytics-vidhya/revealing-bart-a-denoising-objective-for-pretraining-c6e8f8009564)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KArJKXQchHNy"
   },
   "source": [
    "# 6) Context Guided BERT\n",
    "\n",
    "[CG-BERT REPO](https://github.com/frankaging/Quasi-Attention-ABSA/blob/main/code/model/CGBERT.py)\n",
    "\n",
    "```python\n",
    "# number of context classes\n",
    "context_id_map_sentihood = [\n",
    "    'location - 1 - general',\n",
    "    'location - 1 - price',\n",
    "    'location - 1 - safety',\n",
    "    'location - 1 - transit location',\n",
    "    'location - 2 - general',\n",
    "    'location - 2 - price',\n",
    "    'location - 2 - safety',\n",
    "    'location - 2 - transit location'\n",
    "]\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tWeg7wbchJWk"
   },
   "outputs": [],
   "source": [
    "class BERTLayerNorm(nn.Module):\n",
    "    def __init__(self, config, variance_epsilon=1e-12):\n",
    "        super(BERTLayerNorm, self).__init__()\n",
    "        self.gamma = nn.Parameter(torch.ones(config.hidden_size))\n",
    "        self.beta = nn.Parameter(torch.zeros(config.hidden_size))\n",
    "        self.variance_epsilon = variance_epsilon\n",
    "\n",
    "    def forward(self, x):\n",
    "        u = x.mean(-1, keepdim=True)\n",
    "        s = (x - u).pow(2).mean(-1, keepdim=True)\n",
    "        x = (x - u) / torch.sqrt(s + self.variance_epsilon)\n",
    "        return self.gamma * x + self.beta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i0-A5oWPha5j"
   },
   "outputs": [],
   "source": [
    "class ContextBERTSelfAttention(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(ContextBERTSelfAttention, self).__init__()\n",
    "        self.num_attention_heads = config.num_attention_heads\n",
    "        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)\n",
    "        self.all_head_size = self.num_attention_heads * self.attention_head_size\n",
    "\n",
    "        self.query = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "        self.key = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "        self.value = nn.Linear(config.hidden_size, self.all_head_size)\n",
    "        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)\n",
    "\n",
    "        # learnable context integration factors\n",
    "        # enforce initialization to zero as to leave the pretrain model\n",
    "        # unperturbed in the beginning\n",
    "        self.context_for_q = nn.Linear(self.attention_head_size, self.attention_head_size)\n",
    "        self.context_for_k = nn.Linear(self.attention_head_size, self.attention_head_size)\n",
    "\n",
    "        self.lambda_q_context_layer = nn.Linear(self.attention_head_size, 1, bias=False)\n",
    "        self.lambda_q_query_layer = nn.Linear(self.attention_head_size, 1, bias=False)\n",
    "        self.lambda_k_context_layer = nn.Linear(self.attention_head_size, 1, bias=False)\n",
    "        self.lambda_k_key_layer = nn.Linear(self.attention_head_size, 1, bias=False)\n",
    "\n",
    "        # zero-centered activation function, specifically for re-arch fine tunning\n",
    "        self.lambda_sig = nn.Sigmoid()\n",
    "        self.quasi_act = nn.Sigmoid()\n",
    "\n",
    "    def transpose_for_scores(self, x):\n",
    "        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)\n",
    "        x = x.view(*new_x_shape)\n",
    "        return x.permute(0, 2, 1, 3)\n",
    "\n",
    "    def forward(self, hidden_states, attention_mask,\n",
    "                # optional parameters for saving context information\n",
    "                device=None, context_embedded=None):\n",
    "\n",
    "        # (m, seq_len, hidden_dim)\n",
    "        mixed_query_layer = self.query(hidden_states)\n",
    "        mixed_key_layer = self.key(hidden_states)\n",
    "        mixed_value_layer = self.value(hidden_states)\n",
    "\n",
    "        # (m, num_head, seq_len, head_dim)\n",
    "        mixed_query_layer = self.transpose_for_scores(mixed_query_layer)\n",
    "        mixed_key_layer = self.transpose_for_scores(mixed_key_layer)\n",
    "\n",
    "        # (m, 1, hidden_dim) --> (m, num_head, 1, head_dim)\n",
    "        context_embedded = self.transpose_for_scores(context_embedded)\n",
    "        context_embedded_q = self.context_for_q(context_embedded)\n",
    "\n",
    "        # (m, num_head, 1, head_dim) --> (m, num_head, 1, 1)\n",
    "        lambda_q_context = self.lambda_q_context_layer(context_embedded_q)\n",
    "        # (m, num_head, seq_len, head_dim) --> (m, num_head, seq_len, 1)\n",
    "        lambda_q_query = self.lambda_q_query_layer(mixed_query_layer)\n",
    "        # (m, num_head, seq_len, 1)\n",
    "        lambda_q = lambda_q_context + lambda_q_query\n",
    "        lambda_q = self.lambda_sig(lambda_q)\n",
    "\n",
    "        # Q_context = (1-lambda_Q) * Q + lambda_Q * Context_Q\n",
    "        # K_context = (1-lambda_K) * K + lambda_K * Context_K\n",
    "        # the context is shared and is the same for every head.\n",
    "\n",
    "        # (m, num_head, seq_len, head_dim)\n",
    "        contextualized_query_layer = (1 - lambda_q) * mixed_query_layer + lambda_q * context_embedded_q\n",
    "\n",
    "        # repeat same for key\n",
    "        context_embedded_k = self.context_for_k(context_embedded)\n",
    "        lambda_k_context = self.lambda_k_context_layer(context_embedded_k)\n",
    "        lambda_k_key = self.lambda_k_key_layer(mixed_key_layer)\n",
    "        lambda_k = lambda_k_context + lambda_k_key\n",
    "        lambda_k = self.lambda_sig(lambda_k)\n",
    "\n",
    "        # (m, num_head, seq_len, head_dim)\n",
    "        contextualized_key_layer = (1 - lambda_k) * mixed_key_layer + lambda_k * context_embedded_k\n",
    "\n",
    "        ######################################################################\n",
    "\n",
    "        # (m, num_head, seq_len, seq_len)\n",
    "        attention_scores = torch.matmul(\n",
    "            contextualized_query_layer, contextualized_key_layer.transpose(-1, -2)\n",
    "        )\n",
    "        attention_scores = attention_scores / math.sqrt(self.attention_head_size)\n",
    "        attention_scores = attention_scores + attention_mask\n",
    "        attention_probs = nn.Softmax(dim=-1)(attention_scores)\n",
    "        attention_probs = self.dropout(attention_probs)\n",
    "\n",
    "        # (m, num_head, seq_len, seq_len)\n",
    "        value_layer = self.transpose_for_scores(mixed_value_layer)\n",
    "\n",
    "        # (m, num_head, seq_len, head_dim)\n",
    "        context_layer = torch.matmul(attention_probs, value_layer)\n",
    "\n",
    "        # (m, seq_len, num_head, head_dim)\n",
    "        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()\n",
    "\n",
    "        # (m, seq_len, hidden_dim)\n",
    "        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)\n",
    "        context_layer = context_layer.view(*new_context_layer_shape)\n",
    "\n",
    "        return context_layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0DpFJt7QxRZE"
   },
   "outputs": [],
   "source": [
    "class ContextBERTEncoder(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(ContextBERTEncoder, self).__init__()\n",
    "\n",
    "        deep_context_transform_layer = nn.Linear(\n",
    "            2*config.hidden_size, config.hidden_size\n",
    "        )\n",
    "\n",
    "        self.context_layer = nn.ModuleList([\n",
    "            copy.deepcopy(deep_context_transform_layer) for _ in range(config.num_hidden_layers)\n",
    "        ])\n",
    "\n",
    "        layer = ContextBERTLayer(config)\n",
    "        self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])\n",
    "\n",
    "    def forward(self, hidden_states, attention_mask, device=None, context_embeddings=None):\n",
    "\n",
    "        all_encoder_layers = []\n",
    "        layer_index = 0\n",
    "        for layer_module in self.layer:\n",
    "            deep_context_hidden = torch.cat([context_embeddings, hidden_states], dim=-1)\n",
    "            deep_context_hidden = self.context_layer[layer_index](deep_context_hidden)\n",
    "            deep_context_hidden += context_embeddings\n",
    "\n",
    "            # BERT encoding\n",
    "            hidden_states = layer_module(\n",
    "                hidden_states, attention_mask, device, deep_context_hidden\n",
    "            )\n",
    "            # (n_layer, m, seq_len, hidden_dim)\n",
    "            all_encoder_layers.append(hidden_states)\n",
    "            layer_index += 1\n",
    "\n",
    "        return all_encoder_layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WyT5M1jCsl0d"
   },
   "outputs": [],
   "source": [
    "class ContextBertModel(nn.Module):\n",
    "    def __init__(self, config: BertConfig):\n",
    "        super(ContextBertModel, self).__init__()\n",
    "        self.embeddings = BERTEmbeddings(config)\n",
    "        self.encoder = ContextBERTEncoder(config)\n",
    "        self.pooler = ContextBERTPooler(config)\n",
    "\n",
    "        # context embedding\n",
    "        num_target = 4\n",
    "        num_aspect = 2\n",
    "        self.context_embeddings = nn.Embedding(num_target*num_aspect, config.hidden_size)\n",
    "\n",
    "    def forward(self, input_ids, token_type_ids=None, attention_mask=None,\n",
    "                device=None, context_ids=None):\n",
    "\n",
    "        if attention_mask is None:\n",
    "            attention_mask = torch.ones_like(input_ids)\n",
    "        if token_type_ids is None:\n",
    "            token_type_ids = torch.zeros_like(input_ids)\n",
    "\n",
    "        # [batch_size, 1, 1, from_seq_length]\n",
    "        # broadcast to [batch_size, num_heads, seq_length, seq_length]\n",
    "        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)\n",
    "\n",
    "        # multiply by negative big number to make it ignore during softmax\n",
    "        extended_attention_mask = extended_attention_mask.float()\n",
    "        extended_attention_mask = (1.0 - extended_attention_mask) * -9e9\n",
    "        embedding_output = self.embeddings(input_ids, token_type_ids)\n",
    "\n",
    "        # context embeddings\n",
    "        # [batch_size, 1, context_embedding_dim]\n",
    "        context_embedded = self.context_embeddings(context_ids).squeeze(dim=1)\n",
    "        # [batch_size, seq_len, context_embedding_dim]\n",
    "        seq_len = embedding_output.shape[1]\n",
    "        context_embedding_output = torch.stack(seq_len*[context_embedded], dim=1)\n",
    "\n",
    "        # (n_layer, m, seq_len, hidden_dim)\n",
    "        all_encoder_layers = self.encoder(\n",
    "            embedding_output,\n",
    "            extended_attention_mask,\n",
    "            device,\n",
    "            context_embedding_output\n",
    "        )\n",
    "\n",
    "        sequence_output = all_encoder_layers[-1]\n",
    "        pooled_output = self.pooler(sequence_output, attention_mask)\n",
    "        return pooled_output"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
